"""
双积分小车LQR轨迹跟踪控制演示脚本
"""
import os
import time
import numpy as np
import matplotlib.pyplot as plt

from vehicle_model import DoubleIntegratorVehicle
from lqr_controller import LQRController
from trajectory_generator import TrajectoryGenerator
from simulation import Simulation
from visualization import Visualization
from utils import calculate_errors, set_plot_style

def ensure_output_dir(output_dir='results'):
    """确保输出目录存在"""
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        print(f"创建输出目录: {output_dir}")

def run_demo(trajectory_type, q_pos=2000.0, q_vel=200.0, r_acc=0.01, dt=0.1, sim_time=30.0, output_dir='results'):
    """
    运行单个轨迹类型的演示
    """
    print("\n" + "=" * 60)
    print(f"运行 {trajectory_type} 轨迹跟踪演示")
    print("=" * 60)
    
    # 创建系统组件
    vehicle = DoubleIntegratorVehicle(dt=dt)
    traj_gen = TrajectoryGenerator(dt=dt)
    
    # 获取系统矩阵
    A, B = vehicle.get_state_space_matrices()
    
    # 创建LQR控制器
    Q = np.diag([q_pos, q_pos, q_vel, q_vel])
    R = np.diag([r_acc, r_acc])
    controller = LQRController(A, B, Q, R)
    
    # 创建仿真器
    simulator = Simulation(vehicle, controller, traj_gen)
    
    # 创建可视化工具
    visualizer = Visualization()
    
    # 准备仿真
    time_points = np.arange(0, sim_time, dt)
    
    # 生成参考轨迹
    reference_trajectory = traj_gen.generate_trajectory(trajectory_type, time_points)
    
    # 初始状态设置为轨迹起点
    initial_state = reference_trajectory[0].copy()
    
    # 打印仿真配置
    print(f"仿真配置:")
    print(f"  轨迹类型: {trajectory_type}")
    print(f"  初始状态: {initial_state}")
    print(f"  LQR参数: Q_pos={q_pos}, Q_vel={q_vel}, R_acc={r_acc}")
    
    # 首先单独绘制参考轨迹
    plt.figure(figsize=(10, 8))
    plt.plot(reference_trajectory[:, 0], reference_trajectory[:, 1], 'b-', linewidth=2.5, label='参考轨迹')
    plt.scatter(reference_trajectory[0, 0], reference_trajectory[0, 1], c='g', s=100, marker='o', label='起点')
    plt.scatter(reference_trajectory[-1, 0], reference_trajectory[-1, 1], c='r', s=100, marker='x', label='终点')
    
    # 绘制方向箭头
    step = len(reference_trajectory) // 20  # 约20个箭头
    for i in range(0, len(reference_trajectory), step):
        if i + 5 < len(reference_trajectory):
            plt.arrow(reference_trajectory[i, 0], reference_trajectory[i, 1],
                     (reference_trajectory[i+5, 0] - reference_trajectory[i, 0])/2,
                     (reference_trajectory[i+5, 1] - reference_trajectory[i, 1])/2,
                     head_width=0.4, head_length=0.6, fc='blue', ec='blue', alpha=0.6)
    
    plt.title(f'{trajectory_type}参考轨迹', fontsize=16)
    plt.xlabel('X位置', fontsize=14)
    plt.ylabel('Y位置', fontsize=14)
    plt.grid(True)
    plt.axis('equal')
    plt.legend(fontsize=14)
    plt.savefig(f'{output_dir}/{trajectory_type}_reference.png', dpi=300, bbox_inches='tight')
    plt.show()
    
    # 运行仿真
    print("\n开始仿真...")
    start_time = time.time()
    results = simulator.run_simulation(reference_trajectory, initial_state, time_points)
    elapsed_time = time.time() - start_time
    print(f"仿真完成! 耗时: {elapsed_time:.2f}秒")
    
    # 计算误差
    errors = calculate_errors(reference_trajectory[:len(results['states_history'])], results['states_history'])
    print(f"跟踪误差:")
    print(f"  RMSE: {errors['rmse']:.4f}")
    print(f"  最大误差: {errors['max_error']:.4f}")
    print(f"  平均误差: {errors['mean_error']:.4f}")
    
    # 可视化结果
    output_path = f"{output_dir}/{trajectory_type}_tracking.png"
    visualizer.plot_trajectory(results, errors, output_path)
    
    # 生成动画
    animation_path = f"{output_dir}/{trajectory_type}_animation.gif"
    visualizer.create_animation(results, animation_path)
    
    print(f"\n输出文件:")
    print(f"- 参考轨迹: {output_dir}/{trajectory_type}_reference.png")
    print(f"- 跟踪结果: {output_path}")
    print(f"- 动画: {animation_path}")
    
    return results, errors

def run_all_demos():
    """运行所有轨迹类型的演示"""
    # 设置输出目录
    output_dir = 'results'
    ensure_output_dir(output_dir)
    
    # 设置绘图样式
    set_plot_style()
    
    # 轨迹类型
    trajectory_types = ['circle', 'figure_eight', 'square']
    
    # 存储结果
    all_results = {}
    all_errors = {}
    
    # 运行每种轨迹类型的演示
    for traj_type in trajectory_types:
        results, errors = run_demo(traj_type, output_dir=output_dir)
        all_results[traj_type] = results
        all_errors[traj_type] = errors
    
    # 绘制所有轨迹的比较图
    plt.figure(figsize=(12, 10))
    
    for traj_type in trajectory_types:
        results = all_results[traj_type]
        ref_traj = results['reference_trajectory']
        actual_traj = results['states_history']
        
        # 绘制参考轨迹（半透明）
        plt.plot(ref_traj[:, 0], ref_traj[:, 1], '-', linewidth=1.5, alpha=0.4, label=f'{traj_type}参考')
        
        # 绘制实际轨迹
        plt.plot(actual_traj[:, 0], actual_traj[:, 1], '--', linewidth=2, alpha=0.8, label=f'{traj_type}实际')
    
    plt.title('LQR跟踪控制 - 所有轨迹比较', fontsize=16)
    plt.xlabel('X位置', fontsize=14)
    plt.ylabel('Y位置', fontsize=14)
    plt.grid(True)
    plt.axis('equal')
    plt.legend(fontsize=12)
    plt.savefig(f'{output_dir}/all_trajectories_comparison.png', dpi=300, bbox_inches='tight')
    plt.show()
    
    # 打印误差比较
    print("\n" + "=" * 60)
    print("轨迹跟踪误差比较")
    print("=" * 60)
    print(f"{'轨迹类型':<15} {'RMSE':<10} {'最大误差':<10} {'平均误差':<10}")
    print("-" * 60)
    
    for traj_type in trajectory_types:
        errors = all_errors[traj_type]
        print(f"{traj_type:<15} {errors['rmse']:<10.4f} {errors['max_error']:<10.4f} {errors['mean_error']:<10.4f}")
    
    print("=" * 60)
    
    return all_results, all_errors

if __name__ == "__main__":
    run_all_demos() 